# -*- coding: utf-8 -*-


import pymysql

db_map = {}
key41 = "db_41"
key44 = "db_44"

arg_kwargs41 = {
    'host': "192.168.127.41",
    'port': 3306,
    'user': 'root',
    'password': "kg_base_root",
    'database': "pubmed",
    'charset': 'utf8'
}

arg_kwargs44 = {
    'host': "192.168.7.44",
    'port': 3306,
    'user': 'root',
    'password': "kg_base_root",
    'database': "graph_zhsc",
    'charset': 'utf8'
}


def get_connection(key, arg_kwargs):
    db = pymysql.connections.Connection(**arg_kwargs)
    return db
    # if db_map.get(key) is None:
    #     db = pymysql.connections.Connection(**arg_kwargs)
    #     db_map[key] = db
    #     #print("get db from pymysql, key:{0}".format(key))
    #     return db
    # #print("get db from db_map, key:{0}".format(key))
    # return db_map.get(key)


def get_connection44():
    db = pymysql.connections.Connection(**arg_kwargs44)
    return db
    # if db_map.get(key44) is None:
    #     db = pymysql.connections.Connection(**arg_kwargs44)
    #     db_map[key44] = db
    #     #print("get db from pymysql")
    #     return db
    # return db_map.get(key44)


def execute(sql, data):
    db = get_connection(key41, arg_kwargs41)
    cur = db.cursor()
    try:
        cur.executemany(sql, data)
    except Exception as e:
        print("execute sql error, sql:", sql, e)
        db.rollback()
    else:
        db.commit()
    finally:
        cur.close()
        db.close()


def execute_query44(sql):
    db = get_connection(key44, arg_kwargs44)
    cur = db.cursor()
    try:
        print("execute sql:", sql)
        cur.execute(sql)
        result = cur.fetchall()
        return result
    except Exception as e:
        print("execute query error, sql:<{0}>".format(sql), e)
        db.rollback()
    finally:
        db.commit()
        cur.close()
        db.close()


def execute_query(sql):
    db = get_connection(key41, arg_kwargs41)
    cur = db.cursor()
    try:
        cur.execute(sql)
        result = cur.fetchall()
        return result
    except Exception as e:
        print("execute query error, sql:<{}>", sql, e)
        db.rollback()
    finally:
        db.commit()
        cur.close()
        db.close()


def insert_article_node(data):
    '''
    data=[('数学','张三','机械出版社',78,"2023-06-04","数学书"),('英语','李四','机械出版社',67,"2023-07-04","英语书"),('活着','余华','人民出版社',46,"2023-06-01","富贵的一生")]
    :param data:
    :return:
    '''
    sql = ("insert ignore into article (`pmid`,`doi`,`pii`,`pmc`,`article_title`,`medline_pgn`,`publication_status`,"
           "`language`,`vernacular_title`,`abstract`,`citation`) values (%s, %s, %s,%s, %s, %s, %s,%s, %s, %s,%s)")
    execute(sql, data)


def insert_journal_node(data):
    sql = ("insert ignore into journal (`id`,`issn`,`title`,`iso_abbreviation`,`volume`,`issue`,`pub_date`) values ("
           "%s, %s, %s,%s, %s, %s, %s)")
    execute(sql, data)


def insert_article2journal_relation(data):
    '''
    写入article2journal关系数据
    :param data:
    :return:
    '''
    sql = "insert into rel_article_journal (`article_id`, `journal_id`, `volume_issue`) values (%s, %s, %s)"
    execute(sql, data)


def insert_medline_journal_info_node(data):
    '''
    先查再写
    :param data:
    :return:
    '''
    sql = (
        "insert ignore into medline_journal_info (`id`,`nlm_unique_id`,`country`,`medline_ta`,`issn_linking`) values (%s, "
        "%s, %s,%s, %s)")
    execute(sql, data)


def insert_medline2journal_relation(data):
    sql = "insert into rel_medline_journal (`medline_id`, `journal_id`) values (%s, %s)"
    execute(sql, data)


def insert_medline2article_relation(data):
    sql = "insert into rel_article_medline (`pmid`, `medline_id`) values (%s, %s)"
    execute(sql, data)


def insert_mesh_word_node(data):
    '''
    主题词   先查再写
    :param data:
    :return:
    '''
    sql = "insert ignore into mesh_word (`mesh_word_id`, `descriptor_name`, `qualifier_name`) values (%s, %s, %s)"
    execute(sql, data)


def insert_mesh_word2article_relation(data):
    sql = "insert into rel_mesh_word_article (`article_id`, `mesh_word_id`) values (%s, %s)"
    execute(sql, data)


def insert_chemical_node(data):
    '''
    化学品词   先查再写
    :param data:
    :return:
    '''
    sql = "insert ignore into chemical (`chemical_id`, `name_of_substance`, `registry_number`) values (%s, %s, %s)"
    execute(sql, data)


def insert_article2chemical_relation(data):
    sql = "insert into rel_article_chemical (`article_id`, `chemical_id`) values (%s, %s)"
    execute(sql, data)


def insert_publication_type_node(data):
    '''
    出版类型  先查后写
    :param data:
    :return:
    '''
    sql = "insert ignore into publication_type (`id`, `name`) values (%s, %s)"
    execute(sql, data)


def insert_article2publication_type_relation(data):
    sql = "insert into rel_article_publication_type (`article_id`, `publication_type_id`) values (%s, %s)"
    execute(sql, data)


def insert_author_node(data):
    '''
    写入作者数据
    :param data:
    :return:
    '''
    sql = "insert ignore into author (`author_id`, `name`, `Initials`) values (%s, %s, %s)"
    execute(sql, data)


def insert_article2author_relation(data):
    sql = "insert into rel_article_author (`article_id`, `author_id`) values (%s, %s)"
    execute(sql, data)


def insert_reference_node(data):
    '''
    参考文献  合并至article表
    :param data:
    :return:
    '''
    sql = "insert ignore into reference (`pmid`, `title`, `citation`) values (%s, %s, %s)"
    execute(sql, data)


def insert_article2reference_relation(data):
    sql = "insert into rel_article_reference (`article_id`, `reference_article_id`, `citation`) values (%s, %s, %s)"
    execute(sql, data)


def find_max_id(db_name, table_name):
    # sql = 'SELECT MAX(id) FROM ' + table_name
    sql = "SELECT AUTO_INCREMENT FROM information_schema.`TABLES` WHERE Table_Schema= \"" + db_name + "\" AND table_name =\"" + table_name + "\""
    results = execute_query(sql)
    for row in results:
        print('db:{0} table:{1} max id:{2}'.format(db_name, table_name, row[0]))
        return row[0]
    return 0


def update_pubmed_info(node_count, edge_count, row_id):
    if row_id is not None:
        sql = "update kg_parse_pubmed set node_count={0},edge_count={1},status=2 where id={2}".format(node_count,
                                                                                                      edge_count,
                                                                                                      row_id)
        execute_query44(sql)
        print("update kg_parse_pubmed node_count and edge_count stats")


def update_file_dir(file_dir, row_id):
    if file_dir is not None and file_dir.endswith(".xml"):
        sql = 'update kg_parse_pubmed set dir=\"{0}\" where id={1}'.format(file_dir, str(row_id))
        execute_query44(sql)
        print('extract gz and update dir succeed!')


def clear_table_data(db_name, row_id):
    # print("db name:{0}".format(db_name))
    # sql = "SELECT distinct TABLE_NAME FROM information_schema.`COLUMNS` WHERE TABLE_SCHEMA=\"" + db_name + "\""
    stats(db_name=db_name, row_id=row_id)

    sql = "SELECT distinct TABLE_NAME FROM information_schema.`COLUMNS` WHERE TABLE_SCHEMA='pubmed'"
    results = execute_query(sql)
    for result in results:
        clear_sql = "delete from " + str(result[0])
        execute_query(clear_sql)
        print("table <{0}> has been cleared!".format(str(result[0])))
    print("all table has been cleared!")


def stats(db_name, row_id):
    '''
    统计表中的实体数量和关系数量
    :param row_id: 上传记录的id
    :param db_name:  db name
    :return:
    '''
    sql = "SELECT distinct TABLE_NAME FROM information_schema.`COLUMNS` WHERE TABLE_SCHEMA=\"" + db_name + "\""
    results = execute_query(sql)
    rel_count, node_count = 0, 0
    for result in results:
        count_sql = "select count(*) from " + str(result[0])
        count_result = execute_query(count_sql)
        if str(result[0]).startswith("rel"):
            rel_count += count_result[0][0]
        else:
            node_count += count_result[0][0]

    sql2 = 'update kg_parse_pubmed set node_count={0}, edge_count={1} where id={2}'.format(node_count, rel_count, row_id)
    execute_query44(sql2)
    print("update stats success!")


if __name__ == '__main__':
    # clear_table_data("pubmed")
    #find_max_id('pubmed', 'author')
    stats('pubmed_test', 28)
